Summary

This PR adds ExportFriendlyMultiheadAttention, a custom MultiheadAttention implementation that enables torch.export with dynamic shapes (e.g., variable image H/W).

Problem

When exporting SAM3 with dynamic image dimensions using torch.export, the export fails with:

Could not guard on data-dependent expression Eq(u0*u1, 5184)
Caused by: F.multi_head_attention_forward (nn/functional.py:6475)

This happens because nn.MultiheadAttention calls F.multi_head_attention_forward, which has internal guards on sequence length for:

• Shape validation of attention masks
• Fast-path vs slow-path selection
• Self-attention vs cross-attention detection

When the sequence length is symbolic (e.g., H*W where H/W are dynamic), these guards cannot be statically evaluated, causing export to fail.

Note: The commonly suggested workaround sdpa_kernel([SDPBackend.MATH]) does NOT work for this case because the guard failure happens before scaled_dot_product_attention is called - it occurs in F.multi_head_attention_forward's shape validation code.

Solution

ExportFriendlyMultiheadAttention bypasses F.multi_head_attention_forward entirely by:

1. Manually projecting Q, K, V using the same combined in_proj_weight
2. Calling F.scaled_dot_product_attention directly
3. Avoiding all shape validation guards

Also includes:

• from_nn_mha() classmethod to create from existing nn.MultiheadAttention with weight copying
• replace_mha_with_export_friendly() utility function to recursively replace all MHA modules in a model

Test Results

Replaced 61 MultiheadAttention modules

✓ eager_mode: PASS (model works after replacement)
✓ pt2_export: PASS (with strict=False, dynamic H/W works!)
  Graph has 6660 nodes

Usage

from sam3.model.model_misc import replace_mha_with_export_friendly

# Load model
model = build_sam3_image_model(...)

# Replace all MHA modules before export
num_replaced = replace_mha_with_export_friendly(model, verbose=True)
print(f"Replaced {num_replaced} modules")

# Now export with dynamic shapes works
from torch.export import export
from torch.export.dynamic_shapes import Dim

height_mult = Dim("height_mult", min=1, max=6)
width_mult = Dim("width_mult", min=1, max=6)

exported = export(
    model,
    args=(images, ...),
    dynamic_shapes={"images": {2: 336 * height_mult, 3: 336 * width_mult}},
    strict=False,
)

Related Issues

• Could not gaurd on data-dependent expression for SDPA pytorch/pytorch#170127 (SDPA guard error with dynamic shapes)
• [export] Llama3 export with dynamic shapes fails with constraint violations pytorch/pytorch#124502 (Llama3 dynamic shapes export)